
#include "search_manager.h"

#include <algorithm>
#include <iostream>
#include <cmath>


struct Token {
    Token(std::string &str, int p) : token(str), pos(p) {}

    std::string token;
    int pos;
    PostingList posting;
};

struct Result {
    Result(int id, double s) : docID(id), score(s) {}
    int docID;
    double score;
};

SearchManager::SearchManager(std::string indexDB, std::string docDB)
    : indexDBPath_(indexDB), docDBPath_(docDB)
{}

SearchManager::~SearchManager()
{}

int SearchManager::LoadDB()
{
    indexMgr_.LoadDB(indexDBPath_);
    docManager_.LoadDB(docDBPath_);
    return 0;
}

int SearchManager::Search(std::string query)
{
    std::vector<Token> vec = Split(query);

    if (vec.empty()) {
        std::cout << "query is empty\n";
        return 0;
    }

    std::vector<Result> result;
    SearchDocs(vec, query, result);

    Display(result);

    return 0;
}

void SearchManager::Display(std::vector<Result> &result)
{
    std::sort(result.begin(), result.end(), [](Result &a, Result &b) -> bool
    {
        if (a.score > b.score) {
            return true;
        }
        return false;
    });

    for (int i = 0; i < result.size(); ++i) {
        std::string title = docManager_.GetTitle(result[i].docID);
        std::cout << "file " << i << ", title: " << title << "\n";
    }
}

// std::vector<DocID>
int SearchManager::SearchDocs(std::vector<Token> &tokens, std::string &query, std::vector<Result> &result)
{
    // 获取每个token的倒排列表
    auto it = tokens.begin();
    while (it != tokens.end()) {
        it->posting = indexMgr_.GetPosting(it->token);
        ++it;
    }

    //  按照出现的文档数升序排序
    std::sort(tokens.begin(), tokens.end(), [](Token &a, Token &b) -> bool
    {
        if (a.posting.GetDocCount() < b.posting.GetDocCount()) {
            return true;
        }
        return false;
    });

    std::vector<PostingList::DocIterator> iter;
    for (int i = 0; i < tokens.size(); ++i) {
        iter.push_back(tokens[i].posting.DocBegin());
    }

    while (iter[0] != tokens[0].posting.DocEnd()) {
        int docID = tokens[0].posting.GetDocID(iter[0]);
        int nextID = 0;

        for (int i = 1; i < tokens.size(); ++i) {
            while (iter[i] != tokens[i].posting.DocEnd() &&
            tokens[i].posting.GetDocID(iter[i]) < docID) {
                tokens[i].posting.DocNext(iter[i]);
            }

            // 不可能找到交集
            if (iter[i] == tokens[i].posting.DocEnd()) {
                return 0;
            }

            // 需要增加iter0
            if (tokens[i].posting.GetDocID(iter[i]) > docID) {
                nextID = tokens[i].posting.GetDocID(iter[i]);
                break;
            }
        }

        if (nextID > 0) {
            while (iter[0] != tokens[0].posting.DocEnd() &&
                   tokens[0].posting.GetDocID(iter[0]) < nextID) {
                tokens[0].posting.DocNext(iter[0]);
            }
        } else {
            // 找到doc交集
            int phrase = SearchPhrase(tokens, docID);
            if (phrase > 0) {
                double score = CalcScore(tokens);
                result.emplace_back(docID, score);
            }

            // 处理下一个交集
            tokens[0].posting.DocNext(iter[0]);
        }
    }

    return 0;
}

double SearchManager::CalcScore(std::vector<Token> &tokens)
{
    int doc_total = docManager_.GetDocCount();
    double score = 0;

    auto it = tokens.begin();
    while (it != tokens.end()) {
        double idf = std::log2((double)doc_total / it->posting.GetDocCount());
        score += it->posting.GetOccur() * idf;
        ++it;
    }

    return score;
}

int SearchManager::SearchPhrase(std::vector<Token> &tokens, int docID)
{
    int phrase = 0;
    std::vector<PostingList::PosIterator> iters;
    iters.reserve(tokens.size());

    for (int i = 0; i < tokens.size(); ++i) {
        iters.push_back(tokens[i].posting.PosBegin(docID));
    }

    while (iters[0] != tokens[0].posting.PosEnd(docID)) {
        int pos = tokens[0].posting.GetPos(iters[0]) - tokens[0].pos;
        int nextPos = 0;

        for (int i = 1; i < tokens.size(); ++i) {
            while (iters[i] != tokens[i].posting.PosEnd(docID) &&
                    tokens[i].posting.GetPos(iters[i]) - tokens[i].pos < pos) {
                tokens[i].posting.PosNext(iters[i]);
            }

            // 不可能找到交集
            if (iters[i] == tokens[i].posting.PosEnd(docID)) {
                return phrase;
            }

            if (tokens[i].posting.GetPos(iters[i]) - tokens[i].pos > pos) {
                nextPos = tokens[i].posting.GetPos(iters[i]) - tokens[i].pos;
                break;
            }
        }

        if (nextPos > 0) {
            while (iters[0] != tokens[0].posting.PosEnd(docID) &&
                   tokens[0].posting.GetPos(iters[0]) - tokens[0].pos < pos) {
                tokens[0].posting.PosNext(iters[0]);
            }
        } else {
            // 找到交集
            ++phrase;
            tokens[0].posting.PosNext(iters[0]);
        }
    }

    return phrase;
}

std::vector<Token> SearchManager::Split(std::string &str)
{
    std::vector<Token> tokens;
    int pos = 0;

    while (pos < str.size()) {
        std::string token = "";

        int p = pos;
        while (p < str.size() && !std::isalnum(str[p])) {
            ++p;
        }

        pos = p;
        int n = 0;
        while (p + n < str.size() && std::isalnum(str[p + n])) {
            token.push_back(str[p + n]);
            ++n;
        }

        tokens.emplace_back(token, pos);
        pos += n;
    }

    return tokens;
}

